#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdint.h>
#include <math.h>
#include <string.h>

// 大整数类型（64位整数，适配模运算）
typedef int64_t bigint;

// -------------------------- 复用之前的Montgomery算法实现 --------------------------
// 扩展欧几里得算法
int extend_euclid(bigint a, bigint b, bigint *x, bigint *y, bigint *g) {
    if (a < b) {
        int ret = extend_euclid(b, a, y, x, g);
        return ret;
    }

    bigint *coefficients = NULL;
    int coeff_len = 0;
    int coeff_cap = 4;
    coefficients = (bigint *)malloc(coeff_cap * sizeof(bigint));
    if (!coefficients) {
        fprintf(stderr, "扩展欧几里得：内存分配失败！\n");
        return -1;
    }

    while (b != 0) {
        if (coeff_len >= coeff_cap) {
            coeff_cap *= 2;
            bigint *temp = (bigint *)realloc(coefficients, coeff_cap * sizeof(bigint));
            if (!temp) {
                free(coefficients);
                fprintf(stderr, "扩展欧几里得：内存分配失败！\n");
                return -1;
            }
            coefficients = temp;
        }
        coefficients[coeff_len++] = a / b;
        bigint temp = b;
        b = a % b;
        a = temp;
    }

    *x = 1;
    *y = 0;
    *g = a;
    for (int i = coeff_len - 1; i >= 0; i--) {
        bigint new_x = *y;
        bigint new_y = *x - (*y) * coefficients[i];
        *x = new_x;
        *y = new_y;
    }

    free(coefficients);
    return 0;
}

// 丢番图方程求解
int diophantine(bigint a, bigint b, bigint c,
                bigint *x0, bigint *y0, bigint *xk, bigint *yk) {
    bigint x_g, y_g, g;
    int ret = extend_euclid(a, b, &x_g, &y_g, &g);
    if (ret != 0) return -2;

    if (c % g != 0) {
        fprintf(stderr, "丢番图方程：无解！\n");
        return -1;
    }

    bigint k = c / g;
    *x0 = x_g * k;
    *y0 = y_g * k;
    *xk = b / g;
    *yk = -a / g;
    return 0;
}

// Montgomery结构体
typedef struct {
    bigint r;
    bigint n;
    bigint r_;
    bigint n_;
} Montgomery;

// Montgomery初始化
int montgomery_init(Montgomery *mont, bigint r, bigint n) {
    mont->r = r;
    mont->n = n;
    bigint x0, y0, xk, yk;
    int ret = diophantine(r, n, 1, &x0, &y0, &xk, &yk);
    if (ret != 0) return -1;

    mont->r_ = x0;
    mont->n_ = y0;
    while (mont->r_ < 0) {
        if (xk > 0) {
            mont->r_ += xk;
            mont->n_ += yk;
        } else {
            mont->r_ -= xk;
            mont->n_ -= yk;
        }
    }
    return 0;
}

// Montgomery约减
bigint montgomery_reduce(const Montgomery *mont, bigint x) {
    bigint r = mont->r;
    bigint n = mont->n;
    bigint n_ = mont->n_;
    bigint q = ((x % r) * n_) % r;
    bigint a = (x - q * n) / r;
    while (a < 0) a += n;
    return a;
}

// Montgomery域转换
bigint montgomery_present(const Montgomery *mont, bigint x) {
    return (x * mont->r) % mont->n;
}

// Montgomery模幂（核心函数）
bigint montgomery_power(bigint a, bigint e, bigint n) {
    if (n == 1) return 0;
    bigint r = 2;
    while (r <= n) r <<= 1;

    Montgomery mont;
    if (montgomery_init(&mont, r, n) != 0) {
        fprintf(stderr, "Montgomery初始化失败！\n");
        return -1;
    }

    bigint result = montgomery_present(&mont, 1);
    bigint p = montgomery_present(&mont, a % n);
    bigint power = 1;

    while (power <= e) {
        if ((power & e) != 0) {
            result = montgomery_reduce(&mont, result * p);
        }
        p = montgomery_reduce(&mont, p * p);
        power <<= 1;
    }

    return montgomery_reduce(&mont, result);
}

// -------------------------- 大步小步算法实现 --------------------------
/**
 * 哈希表节点（存储 <baby值, q> 键值对）
 * 用于快速查找baby步骤的结果
 */
typedef struct HashNode {
    bigint baby_val;   // 键：baby = b*a^q mod m
    int q;             // 值：对应的q
    struct HashNode *next; // 链表解决哈希冲突
} HashNode;

/**
 * 哈希表初始化
 * @param size: 哈希表大小（建议为素数，此处用n+1）
 * @return: 哈希表指针
 */
HashNode **hash_table_create(int size) {
    HashNode **table = (HashNode **)calloc(size, sizeof(HashNode *));
    if (!table) {
        fprintf(stderr, "哈希表：内存分配失败！\n");
        exit(EXIT_FAILURE);
    }
    return table;
}

/**
 * 哈希函数（简单取模）
 */
static int hash_func(bigint val, int size) {
    return (val % size + size) % size; // 确保非负
}

/**
 * 哈希表插入键值对
 */
void hash_table_insert(HashNode **table, int size, bigint baby_val, int q) {
    int idx = hash_func(baby_val, size);
    HashNode *new_node = (HashNode *)malloc(sizeof(HashNode));
    new_node->baby_val = baby_val;
    new_node->q = q;
    new_node->next = table[idx]; // 头插法
    table[idx] = new_node;
}

/**
 * 哈希表查找baby_val对应的q
 * @return: 找到返回q，未找到返回-1
 */
int hash_table_find(HashNode **table, int size, bigint baby_val) {
    int idx = hash_func(baby_val, size);
    HashNode *curr = table[idx];
    while (curr) {
        if (curr->baby_val == baby_val) {
            return curr->q;
        }
        curr = curr->next;
    }
    return -1;
}

/**
 * 哈希表销毁
 */
void hash_table_destroy(HashNode **table, int size) {
    for (int i = 0; i < size; i++) {
        HashNode *curr = table[i];
        while (curr) {
            HashNode *temp = curr;
            curr = curr->next;
            free(temp);
        }
    }
    free(table);
}

/**
 * 整数二分法求平方根：计算 floor(sqrt(x))，x为非负整数
 * @param x: 待求平方根的数（x >=0）
 * @return: 最大的整数s，满足 s^2 <= x < (s+1)^2
 */
bigint integer_sqrt(bigint x) {
    if (x < 0) {
        fprintf(stderr, "错误：平方根参数不能为负数！\n");
        return -1;
    }
    if (x == 0 || x == 1) return x;

    bigint low = 1;
    bigint high = x / 2; // 平方根最大不超过x/2（x>=2时）
    bigint ans = 0;

    while (low <= high) {
        bigint mid = low + (high - low) / 2; // 避免溢出
        bigint mid_sq = mid * mid;

        if (mid_sq == x) {
            return mid; // 完全平方数，直接返回
        } else if (mid_sq < x) {
            ans = mid;   // 记录当前可能的最大解
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }

    return ans; // 返回floor(sqrt(x))
}

/**
 * 大步小步算法：求解 a^x ≡ b (mod m) 的最小非负整数x
 * @param a: 底数
 * @param b: 结果
 * @param m: 模数
 * @return: 解x；无解返回-1
 */
bigint baby_step_giant_step(bigint a, bigint b, bigint m) {
    // 预处理：a和b对m取模
    a = a % m;
    b = b % m;

    // 特殊情况：b=1时，x=0
    if (b == 1) return 0;

    // 计算n = floor(sqrt(m)) + 1
    int n = (int)integer_sqrt(m) + 1;

    // 第一步：Baby Step（存储 b*a^q mod m → q）
    HashNode **giants_table = hash_table_create(n + 1); // 哈希表大小为n+1
    for (int q = 0; q <= n; q++) {
        // 计算 baby = b * a^q mod m
        bigint a_q = montgomery_power(a, q, m); // a^q mod m
        bigint baby = (b * a_q) % m;
        hash_table_insert(giants_table, n + 1, baby, q);
    }

    // 第二步：Giant Step（计算 a^(p*n) mod m，查找是否在Baby Step中）
    bigint a_n = montgomery_power(a, n, m); // a^n mod m
    for (int p = 1; p <= n; p++) {
        // 计算 giant = a^(p*n) mod m
        bigint giant = montgomery_power(a_n, p, m);
        // 查找giant对应的q
        int q = hash_table_find(giants_table, n + 1, giant);
        if (q != -1) {
            // 找到解：x = p*n - q
            bigint x = (bigint)p * n - q;
            hash_table_destroy(giants_table, n + 1);
            return x;
        }
    }

    // 无解
    hash_table_destroy(giants_table, n + 1);
    fprintf(stderr, "RuntimeError: 无解\n");
    return -1;
}

// -------------------------- 测试主函数 --------------------------
int main() {
    // 测试用例：对应Python的 a=9, x=99, m=7
    bigint a = 9, x_true = 99, m = 7;

    // 计算 b = a^x_true mod m
    bigint b = montgomery_power(a, x_true, m);
    printf("原始参数：a=%lld, x=%lld, m=%lld\n", a, x_true, m);
    printf("计算 b = a^x mod m = %lld\n", b);

    // 用大步小步算法求解x
    bigint x_solve = baby_step_giant_step(a, b, m);
    if (x_solve != -1) {
        printf("求解得到 x = %lld\n", x_solve);
        // 验证：计算 a^x_solve mod m，是否等于b
        bigint b_verify = montgomery_power(a, x_solve, m);
        printf("验证：a^x_solve mod m = %lld, 与b=%lld %s\n",
               b_verify, b, (b_verify == b) ? "一致" : "不一致");
    }

    // 额外测试用例：a=3, b=14, m=17（3^777 ≡14 mod17，求解x=777）
    printf("\n额外测试用例：\n");
    bigint a2 = 3, b2 = 14, m2 = 17;
    bigint x2 = baby_step_giant_step(a2, b2, m2);
    if (x2 != -1) {
        printf("求解 3^x ≡14 mod17 的解：x=%lld\n", x2);
        bigint b2_verify = montgomery_power(a2, x2, m2);
        printf("验证：3^%lld mod17 = %lld, 与b=%lld %s\n",
               x2, b2_verify, b2, (b2_verify == b2) ? "一致" : "不一致");
    }

    return 0;
}